# Copyright (c) Facebook, Inc. and its affiliates.
import logging
from copy import deepcopy
from typing import Callable, Dict, List, Optional, Tuple, Union

import fvcore.nn.weight_init as weight_init

import torch
from torch import nn
from torch.nn import functional as F

import numpy as np

from detectron2.config import configurable
from detectron2.layers import Conv2d, ShapeSpec, get_norm
from detectron2.modeling import SEM_SEG_HEADS_REGISTRY
from detectron2.structures import Boxes, ImageList, Instances, BitMasks
from detectron2.modeling.postprocessing import sem_seg_postprocess
from detectron2.utils.memory import retry_if_cuda_oom
from detectron2.modeling.backbone.resnet import BasicBlock

from ..transformer_decoder.maskformer_transformer_decoder import build_transformer_decoder
from ..pixel_decoder.fpn import build_pixel_decoder
from ..transformer_decoder.instance_retrieval_decoder import InstanceRetrievalDecoder

@SEM_SEG_HEADS_REGISTRY.register()
class GirCGC2FormerHead(nn.Module):
    _version = 2

    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
        version = local_metadata.get("version", None)
        if version is None or version < 2:
            # Do not warn if train from scratch
            scratch = True
            logger = logging.getLogger(__name__)
            for k in list(state_dict.keys()):
                newk = k
                if "sem_seg_head" in k and not k.startswith(prefix + "predictor"):
                    newk = k.replace(prefix, prefix + "pixel_decoder.")
                    # logger.debug(f"{k} ==> {newk}")
                if newk != k:
                    state_dict[newk] = state_dict[k]
                    del state_dict[k]
                    scratch = False

            if not scratch:
                logger.warning(
                    f"Weight format of {self.__class__.__name__} have changed! "
                    "Please upgrade your models. Applying automatic conversion now ..."
                )

    @configurable
    def __init__(
        self,
        input_shape: Dict[str, ShapeSpec],
        *,
        num_classes: int,
        pixel_decoder: nn.Module,
        loss_weight: float = 1.0,
        ignore_value: int = -1,
        # extra parameters
        transformer_predictor: nn.Module,
        transformer_in_feature: str,
        train_retrieval_on: bool = True,
        train_mask_on: bool = True,
        test_retrieval_on: bool = True,
        test_mask_on: bool = True
    ):
        """
        NOTE: this interface is experimental.
        Args:
            input_shape: shapes (channels and stride) of the input features
            num_classes: number of classes to predict
            pixel_decoder: the pixel decoder module
            loss_weight: loss weight
            ignore_value: category id to be ignored during training.
            transformer_predictor: the transformer decoder that makes prediction
            transformer_in_feature: input feature name to the transformer_predictor
        """
        super().__init__()
        input_shape = sorted(input_shape.items(), key=lambda x: x[1].stride)
        self.in_features = [k for k, v in input_shape]
        feature_strides = [v.stride for k, v in input_shape]
        feature_channels = [v.channels for k, v in input_shape]

        self.ignore_value = ignore_value
        self.common_stride = 4
        self.loss_weight = loss_weight

        self.pixel_decoder = pixel_decoder

        #decoder
        self.predictor = transformer_predictor
        self.transformer_in_feature = transformer_in_feature        
        self.num_classes = num_classes
        self.test_mask_on = test_mask_on          
        self.test_retrieval_on = test_retrieval_on
        
        self.train_mask_on = train_mask_on          
        self.train_retrieval_on = train_retrieval_on
        
        self.channels_task1 = 192
        self.channels_task2 = 64
        
        
        if self.train_retrieval_on:
            #instance decoder
            self.instance_decoder = InstanceRetrievalDecoder()
        

            self.min_box_wh = 8.0
            shared_nums = 128
            # self.gate_task1 = nn.Linear(shared_nums, shared_nums)
            # self.gate_task2 = nn.Linear(shared_nums, shared_nums)
            self.gate_task1 = BasicBlock(self.channels_task1, self.channels_task1, stride=1)
            self.gate_task2 = BasicBlock(self.channels_task2, self.channels_task2, stride=1)

            # self.gate_task1_norm = nn.LayerNorm(self.channels_task1)
            # self.gate_task2_norm = nn.LayerNorm(self.channels_task2)
            
            self.gate_task2_proj = nn.Linear(256, 256)
            self.gate_task1_proj = nn.Linear(256, 256)
            
        self._shift_box_flag = False
        
        
    @classmethod
    def from_config(cls, cfg, input_shape: Dict[str, ShapeSpec]):
        # figure out in_channels to transformer predictor
        if cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE == "transformer_encoder":
            transformer_predictor_in_channels = cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM
        elif cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE == "pixel_embedding":
            transformer_predictor_in_channels = cfg.MODEL.SEM_SEG_HEAD.MASK_DIM
        elif cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE == "multi_scale_pixel_decoder":  # for maskformer2
            transformer_predictor_in_channels = cfg.MODEL.SEM_SEG_HEAD.CONVS_DIM
        else:
            transformer_predictor_in_channels = input_shape[cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE].channels
        return {
            "input_shape": {
                k: v for k, v in input_shape.items() if k in cfg.MODEL.SEM_SEG_HEAD.IN_FEATURES
            },
            "ignore_value": cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE,
            "num_classes": cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES,
            "pixel_decoder": build_pixel_decoder(cfg, input_shape),
            "loss_weight": cfg.MODEL.SEM_SEG_HEAD.LOSS_WEIGHT,
            "transformer_in_feature": cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE,
            "transformer_predictor": build_transformer_decoder(
                cfg,
                transformer_predictor_in_channels,
                mask_classification=True,
            ),
            "train_retrieval_on": cfg.MODEL.SEM_SEG_HEAD.TRAIN_RETRIEVAL_ON,
            "train_mask_on": cfg.MODEL.SEM_SEG_HEAD.TRAIN_MASK_ON,
            "test_retrieval_on": cfg.MODEL.MASK_FORMER.TEST.RETRIEVAL_ON,
            "test_mask_on": cfg.MODEL.MASK_FORMER.TEST.INSTANCE_ON
            
        }
    def getBoxesinTargets(self,targets):
        
        b_tensor = torch.tensor([],dtype=torch.int)
        ids_tensor = torch.tensor([],dtype=torch.int)
        boxes_tensor = torch.tensor([],dtype=torch.int)
        
        bs = len(targets)
        for b in range(bs):
            tgt_masks = targets[b]["masks"]
            tgt_ids = targets[b]["gt_ids"]
            boxes = BitMasks(tgt_masks).get_bounding_boxes()
            
            valid_box = boxes.nonempty(threshold=self.min_box_wh)
            
            boxes = boxes[valid_box]
            tgt_ids = tgt_ids[valid_box].to(ids_tensor)

            if(len(boxes) < 1):
                continue
            boxes4 = np.ceil(boxes.tensor / 4 )
            boxes4= boxes4.to(torch.int)
            
            teb = torch.full((len(boxes),), b)
            boxes_tensor = torch.cat((boxes_tensor,boxes4),dim = 0)
            ids_tensor = torch.cat((ids_tensor,tgt_ids),dim=0)
            b_tensor = torch.cat((b_tensor,teb),dim=0)
        
        return boxes_tensor,b_tensor,ids_tensor
    def expand_box(self,pred_boxes,max_w,max_h):
        boxes_ts  = pred_boxes
        ww = boxes_ts[:,2] - boxes_ts[:,0]
        hh = boxes_ts[:,3] - boxes_ts[:,1]
        dw = torch.sqrt(ww)//2
        dh = torch.sqrt(hh)//2
        
        ts0 = boxes_ts[:,0]
        ts1 = boxes_ts[:,1]
        ts2 = boxes_ts[:,2]
        ts3 = boxes_ts[:,3]
        ts0 = torch.where(ts0 > dw,ts0-dw,ts0)
        ts1 = torch.where(ts1 > dh,ts1-dh,ts1)
        ts2 = torch.where(ts2+dw < max_w,ts2+dw,ts2)
        ts3 = torch.where(ts3+dh < max_h,ts3+dh,ts3)
        
        boxes_ts_2 = boxes_ts.clone()
        boxes_ts[:,0] = ts0
        boxes_ts[:,1] = ts1
        boxes_ts[:,2] = ts2
        boxes_ts[:,3] = ts3
        return boxes_ts
            
    def getBoxesinTargetsWithShift(self,targets):
        
        b_tensor = torch.tensor([],dtype=torch.int)
        ids_tensor = torch.tensor([],dtype=torch.int)
        boxes_tensor = torch.tensor([],dtype=torch.int)
        boxes_tensor_2 = torch.tensor([],dtype=torch.int)
        _,imw,imh = targets[0]['masks'].shape
        bs = len(targets)
        for b in range(bs):
            tgt_masks = targets[b]["masks"]
            tgt_ids = targets[b]["gt_ids"]
            boxes = BitMasks(tgt_masks).get_bounding_boxes()
            
            valid_box = boxes.nonempty(threshold=self.min_box_wh)
            
            boxes = boxes[valid_box]
            tgt_ids = tgt_ids[valid_box].to(ids_tensor)
            if(len(boxes) < 1):
                continue
            
            boxes_ts  = boxes.tensor
            ww = boxes_ts[:,2] - boxes_ts[:,0]
            hh = boxes_ts[:,3] - boxes_ts[:,1]
            
            dw = torch.sqrt(ww)
            dh = torch.sqrt(hh)
            sn = torch.rand(size=(dw.shape[0],2)) * 2 - 1
            dw1 = dw * sn[:,0]
            dh1 = dh * sn[:,1]
            dwh = torch.zeros(size=(dw.shape[0],4))
            dwh[:,0] = dw1
            dwh[:,1] = dh1
            dwh[:,2] = dw1
            dwh[:,3] = dh1
            
            boxes_ts_2= boxes_ts + dwh
            boxes_ts_3= boxes_ts - dwh

            for inds in range(boxes_ts.shape[0]):
                if (boxes_ts_2[inds,0] < 0 or boxes_ts_2[inds,1] < 0 or boxes_ts_2[inds,2] > imw or boxes_ts_2[inds,3] > imh) :
                    boxes_ts_2[inds] = boxes_ts[inds]
                if (boxes_ts_3[inds,0] < 0 or boxes_ts_3[inds,1] < 0 or boxes_ts_3[inds,2] > imw or boxes_ts_3[inds,3] > imh) :
                    boxes_ts_3[inds] = boxes_ts[inds]
                    
            # boxes_ts_2 = torch.cat((boxes_ts_2,boxes_ts_3),dim = 0)
            boxes2c = np.ceil(boxes_ts_2 / 4 )
            boxes2c= boxes2c.to(torch.int)

            boxes3c = np.ceil(boxes_ts_3 / 4 )
            boxes3c= boxes3c.to(torch.int)
            
            boxes_tensor = torch.cat((boxes_tensor,boxes2c),dim = 0)
            boxes_tensor_2 = torch.cat((boxes_tensor,boxes3c),dim = 0)
            
            ids_tensor = torch.cat((ids_tensor,tgt_ids),dim=0)
            teb = torch.full((len(boxes),), b)
            b_tensor = torch.cat((b_tensor,teb),dim=0)
        
        return boxes_tensor,boxes_tensor_2,b_tensor,ids_tensor

    
    def gate_process(self,mask_features):
        
        def __gen_2f(feature,gfn1,gfn2):
            f1 = feature[:,0:self.channels_task1,:,:]
            f2 = feature[:,self.channels_task1:256,:,:]
            f11 = f1.detach()
            f22 = f2.detach()
            
            f22 = self.gate_task2(f22)
            f11 = self.gate_task1(f11)
            f22 = F.layer_norm(f22.permute(0,2,3,1),[self.channels_task2]).permute(0,3,1,2)
            f11 = F.layer_norm(f11.permute(0,2,3,1),[self.channels_task1]).permute(0,3,1,2)
            # f22 = self.gate_task2_norm(f22)
            # f11 = self.gate_task1_norm(f11)
            
            f3_1 = torch.cat((f1,f22),dim=1)
            f3_3 = torch.cat((f11,f2),dim=1)
            return f3_1,f3_3
            
        mf1,mf2 = __gen_2f(mask_features,self.gate_task1,self.gate_task2)
        
        
        def __proj(feature,g_fn):
            #feature B*CGATE*H*W
            b, cg, h, w = feature.shape
            feature = feature.flatten(2).permute(0, 2, 1)
            feature = g_fn(feature)
            cg2 = g_fn.out_features
            feature = feature.permute(0, 2, 1).view(b, cg2, h, w)  
            return feature
        mf1 = __proj(mf1,self.gate_task1_proj)
        mf2 = __proj(mf2,self.gate_task2_proj)
        return mf1,mf2
          
    
    def forward(self, features, targets=None):
        #mask no use
        mask = None
        #mask_features encoder的最后一次层输出
        #multi_scale_features encoder的多层输出 3层
        '''
        SRC IMAGE 10*512*512
        features:res2 10, 256, 128, 128
        features:res2 10, 512, 64, 64
        features:res2 10, 1024, 32, 32
        features:res2 10, 2048, 16, 16
        mask_features :10, 256, 128, 128
        transformer_encoder_features :10, 256, 16, 16
        multi_scale_features[0]:10, 256, 16, 16
        multi_scale_features[1]:10, 256, 32, 32
        multi_scale_features[2]:10, 256, 64, 64
        '''        
        mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
        if self.train_retrieval_on :
            mf1,mf2 = self.gate_process(mask_features)
        else:
            mf1 = mask_features
        if self.transformer_in_feature == "multi_scale_pixel_decoder":
            predictions = {}
            if self.train_mask_on:
                predictions = self.predictor(multi_scale_features,mf1, mask)
            #predictions = self.predictor(multi_scale_features, mask_features, mask)
            if self.train_retrieval_on:
                if self._shift_box_flag:
                    boxes_tensor,boxes_tensor_2,b_tensor,ids_tensor = self.getBoxesinTargetsWithShift(targets)
                    if len(ids_tensor) > 0:
                        perd_instances = self.instance_decoder(mf2,boxes_tensor,b_tensor)
                        perd_instances_2 = self.instance_decoder(mf2,boxes_tensor_2,b_tensor)
                        predictions['instance_perd_feats'] = perd_instances
                        predictions['instance_tgt_ids'] = ids_tensor
                        predictions['instance_perd_feats_2'] = perd_instances_2
                else:
                    boxes_tensor,b_tensor,ids_tensor = self.getBoxesinTargets(targets)
                    if len(ids_tensor) > 0:
                        perd_instances = self.instance_decoder(mf2,boxes_tensor,b_tensor)
                        predictions['instance_perd_feats'] = perd_instances
                        predictions['instance_tgt_ids'] = ids_tensor   
                
            # b_id_box = torch.cat((b_tensor.unsqueeze(dim=1),ids_tensor.unsqueeze(dim=1),boxes_tensor),dim=1)
            # predictions['instance_tgt_bidbox'] = b_id_box
            
        else:
            if self.transformer_in_feature == "transformer_encoder":
                assert (
                    transformer_encoder_features is not None
                ), "Please use the TransformerEncoderPixelDecoder."
                predictions = self.predictor(transformer_encoder_features, mask_features, mask)
            elif self.transformer_in_feature == "pixel_embedding":
                predictions = self.predictor(mask_features, mask_features, mask)
            else:
                predictions = self.predictor(features[self.transformer_in_feature], mask_features, mask)
        return predictions
        
    def forward_eval(self,features,batched_inputs, images):
        #
        mask = None
        mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
        if self.train_retrieval_on:
            mf1,mf2 = self.gate_process(mask_features)
        else:
            mf1 = mask_features
        if self.transformer_in_feature == "multi_scale_pixel_decoder":
            predictions = self.predictor(multi_scale_features, mf1, mask)
        else:
            if self.transformer_in_feature == "transformer_encoder":
                assert (
                    transformer_encoder_features is not None
                ), "Please use the TransformerEncoderPixelDecoder."
                predictions = self.predictor(transformer_encoder_features, mask_features, mask)
            elif self.transformer_in_feature == "pixel_embedding":
                predictions = self.predictor(mask_features, mask_features, mask)
            else:
                predictions = self.predictor(features[self.transformer_in_feature], mask_features, mask)        
        
        mask_cls_results = predictions["pred_logits"]
        mask_pred_results = predictions["pred_masks"]
        # upsample masks
        mask_pred_results = F.interpolate(
            mask_pred_results,
            size=(images.tensor.shape[-2], images.tensor.shape[-1]),
            mode="bilinear",
            align_corners=False,
        )
        
        processed_results = []
        b_tensor = torch.tensor([],dtype=torch.int)
        boxes_tensor = torch.tensor([],dtype=torch.int)
        
        batch_index = -1
        for mask_cls_result, mask_pred_result, input_per_image, image_size in zip(
            mask_cls_results, mask_pred_results, batched_inputs, images.image_sizes
        ):
            batch_index += 1
            height = input_per_image.get("height", image_size[0])
            width = input_per_image.get("width", image_size[1])

            processed_results.append({})


            if self.test_retrieval_on:
                ### 计算INSTANCE MIX FEATURE
                pred_boxes = BitMasks(mask_pred_result > 0).get_bounding_boxes()

                ###
            # 这个函数的功能是将mask_pred_result 调整为原始图像输入的H和W
            mask_pred_result = retry_if_cuda_oom(sem_seg_postprocess)(
                mask_pred_result, image_size, height, width
            )
            instance_r,resort_topk_indices = retry_if_cuda_oom(self.instance_mask_inference)(mask_cls_result, mask_pred_result)
            
            if self.test_retrieval_on:
                resort_topk_indices = resort_topk_indices.to(pred_boxes.tensor.device)
                pred_boxes.tensor = pred_boxes.tensor[resort_topk_indices]
                
                valid_box_index = pred_boxes.nonempty(threshold=self.min_box_wh-2.0)
                valid_boxes = pred_boxes[valid_box_index]
                pred_feature = torch.zeros((pred_boxes.tensor.shape[0],self.instance_decoder.out_feature_dim))
                
                if(len(valid_boxes) > 0):
                    valid_boxes = np.ceil(valid_boxes.tensor / 4 )
                    valid_boxes = self.expand_box(valid_boxes,112,112)
                    valid_boxes= valid_boxes.to(torch.int)
                    perd_mix_feats = self.instance_decoder.inference_one(mf2[[batch_index]],valid_boxes)
                    # pred_feature = pred_feature.to(perd_mix_feats)
                    perd_mix_feats = perd_mix_feats.to(pred_feature)
                    pred_feature[valid_box_index] = perd_mix_feats
                instance_r.pred_boxes_ins_mask = valid_box_index
                instance_r.pred_boxes_ins_feature = pred_feature
                
            processed_results[-1]["instances"] = instance_r
            
        return processed_results
   
    def instance_mask_inference(self, mask_cls, mask_pred):
        
        test_topk_per_image = 100
        num_queries = 100
        # mask_pred is already processed to have the same shape as original input
        image_size = mask_pred.shape[-2:]

        # [Q, K]
        scores = F.softmax(mask_cls, dim=-1)[:, :-1]
        labels = torch.arange(self.num_classes).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1)
        labels = labels.to(scores)
        # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False)
        scores_per_image, topk_indices = scores.flatten(0, 1).topk(test_topk_per_image, sorted=False)
        labels_per_image = labels[topk_indices]
        
        topk_indices = topk_indices // self.num_classes
        #mask_pred = mask_pred.unsqueeze(1).repeat(1, self.sem_seg_head.num_classes, 1).flatten(0, 1)
        mask_pred = mask_pred[topk_indices]

        result = Instances(image_size)
        # mask (before sigmoid)
        result.pred_masks = (mask_pred > 0).float()
        
        #result.pred_boxes = Boxes(torch.zeros(mask_pred.size(0), 4))
        # Uncomment the following to get boxes from masks (this is slow)
        result.pred_boxes = BitMasks(mask_pred > 0).get_bounding_boxes()

        # calculate average mask prob
        mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6)
        result.scores = scores_per_image * mask_scores_per_image
        result.pred_classes = labels_per_image.to(torch.int)
        return result,topk_indices

#region noclass
    def forward_eval_noclass(self,features,batched_inputs, images):
        #
        mask = None
        mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
        if self.transformer_in_feature == "multi_scale_pixel_decoder":
            predictions = self.predictor(multi_scale_features, mask_features, mask)
        else:
            if self.transformer_in_feature == "transformer_encoder":
                assert (
                    transformer_encoder_features is not None
                ), "Please use the TransformerEncoderPixelDecoder."
                predictions = self.predictor(transformer_encoder_features, mask_features, mask)
            elif self.transformer_in_feature == "pixel_embedding":
                predictions = self.predictor(mask_features, mask_features, mask)
            else:
                predictions = self.predictor(features[self.transformer_in_feature], mask_features, mask)
                
        
        # mask_cls_results = predictions["pred_logits"]
        
        mask_pred_results = predictions["pred_masks"]
        # upsample masks
        mask_pred_results = F.interpolate(
            mask_pred_results,
            size=(images.tensor.shape[-2], images.tensor.shape[-1]),
            mode="bilinear",
            align_corners=False,
        )
        
        processed_results = []
        b_tensor = torch.tensor([],dtype=torch.int)
        boxes_tensor = torch.tensor([],dtype=torch.int)
        
        batch_index = -1
        # for mask_cls_result, mask_pred_result, input_per_image, image_size in zip(
        #     mask_cls_results, mask_pred_results, batched_inputs, images.image_sizes
        # ):
        for mask_pred_result, input_per_image, image_size in zip(
            mask_pred_results, batched_inputs, images.image_sizes
        ):
            batch_index += 1
            height = input_per_image.get("height", image_size[0])
            width = input_per_image.get("width", image_size[1])
            # height = image_size[0]
            # width = image_size[1]
            processed_results.append({})

            # 这个函数的功能是将mask_pred_result 调整为原始图像输入的H和W
            mask_pred_result_orisize = retry_if_cuda_oom(sem_seg_postprocess)(
                mask_pred_result, image_size, height, width
            )
            # mask_cls_result = mask_cls_result.to(mask_pred_result)
            mask_cls_result = None
            instance_r = retry_if_cuda_oom(self.instance_mask_inference_noclass)(mask_cls_result, mask_pred_result_orisize)
            
            ### 计算INSTANCE MIX FEATURE
            pred_boxes = BitMasks(mask_pred_result > 0).get_bounding_boxes()
            valid_box_index = pred_boxes.nonempty(threshold=8.0)
            valid_boxes = pred_boxes[valid_box_index]
            pred_feature = torch.zeros((pred_boxes.tensor.shape[0],self.instance_decoder.out_feature_dim))
            
            if(len(valid_boxes) > 0):
                valid_boxes = np.ceil(valid_boxes.tensor / 4 )
                valid_boxes= valid_boxes.to(torch.int)
                perd_mix_feats = self.instance_decoder.inference(
                    mask_features,multi_scale_features,batch_index,valid_boxes)
                pred_feature = pred_feature.to(perd_mix_feats)
                pred_feature[valid_box_index] = perd_mix_feats

            instance_r.pred_boxes_ins_mask = valid_box_index
            instance_r.pred_boxes_ins_feature = pred_feature
            ###

            processed_results[-1]["instances"] = instance_r
            
        return processed_results
   
    def instance_mask_inference_noclass(self, mask_cls, mask_pred):
        
        test_topk_per_image = 100
        num_queries = 100
        # mask_pred is already processed to have the same shape as original input
        image_size = mask_pred.shape[-2:]

        # [Q, K]
        # scores = F.softmax(mask_cls, dim=-1)[:, :-1]
        # labels = torch.arange(self.num_classes).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1)
        # labels = labels.to(mask_cls)
        # # scores_per_image, topk_indices = scores.flatten(0, 1).topk(self.num_queries, sorted=False)
        # scores_per_image, topk_indices = scores.flatten(0, 1).topk(test_topk_per_image, sorted=False)
        # labels_per_image = labels[topk_indices]

        labels_per_image = torch.zeros(test_topk_per_image).to(mask_pred)
        
        # topk_indices = topk_indices // self.num_classes
        # mask_pred = mask_pred.unsqueeze(1).repeat(1, self.sem_seg_head.num_classes, 1).flatten(0, 1)
        # mask_pred = mask_pred[topk_indices]

        result = Instances(image_size)
        # mask (before sigmoid)
        result.pred_masks = (mask_pred > 0).float()
        
        #result.pred_boxes = Boxes(torch.zeros(mask_pred.size(0), 4))
        # Uncomment the following to get boxes from masks (this is slow)
        result.pred_boxes = BitMasks(mask_pred > 0).get_bounding_boxes()

        # calculate average mask prob
        mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * result.pred_masks.flatten(1)).sum(1) / (result.pred_masks.flatten(1).sum(1) + 1e-6)
        result.scores = mask_scores_per_image
        # result.scores = scores_per_image * mask_scores_per_image
        result.pred_classes = labels_per_image.to(torch.int)
        return result
#endregion

    #mask no use
    def layers(self, features, mask=None):
        #
        mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
        if self.transformer_in_feature == "multi_scale_pixel_decoder":
            predictions = self.predictor(multi_scale_features, mask_features, mask)
        else:
            if self.transformer_in_feature == "transformer_encoder":
                assert (
                    transformer_encoder_features is not None
                ), "Please use the TransformerEncoderPixelDecoder."
                predictions = self.predictor(transformer_encoder_features, mask_features, mask)
            elif self.transformer_in_feature == "pixel_embedding":
                predictions = self.predictor(mask_features, mask_features, mask)
            else:
                predictions = self.predictor(features[self.transformer_in_feature], mask_features, mask)
        return predictions
